from flask import Flask, request

from src.basic_model.basic_predict_fun import basic_predict
from src.bert.bert_predict import bert_predict
from src.llm.llm_classify import get_zhipuai_res

app = Flask(__name__)
from flask_cors import CORS

CORS(app, resources={r"/predict": {"origins": "http://localhost:3001"}})  # 允许前端域名访问

# 根据模型 ID 进行预测的函数
@app.route('/predict', methods=['POST'])
def predict_api():
    data = request.get_json()
    modelId = data.get('modelId')
    text = data.get('text')
    # print(data)
    result = predict(texts=text, modelId=modelId)
    return result


def predict(texts, task="all", modelId='bert'):
    # 确保文本始终是一个列表形式。
    if isinstance(texts, str):  # 如果传入的参数是一个字符串，则将其转换为列表。
        texts = [texts]

    # 根据“模型ID”，调用相应的函数
    if modelId == 'bert':
        return bert_predict(texts)
    elif modelId == 'base':
        return basic_predict(texts, task)
    elif modelId == 'deepseek':
        return get_zhipuai_res('\n'.join(texts))


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8080)
